{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b31c0a84",
   "metadata": {},
   "source": [
    "## Rerun imports and initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1076c3a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import math\n",
    "from collections import namedtuple\n",
    "from math import cos, sin, tau\n",
    "\n",
    "import numpy as np\n",
    "import rerun as rr  # pip install rerun-sdk\n",
    "import rerun.blueprint as rrb\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bf894a1f",
   "metadata": {},
   "source": [
    "## Helper to create the colored cube\n",
    "\n",
    "This code exists in the `rerun.utilities` package, but is included here for context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f709925e",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "ColorGrid = namedtuple(\"ColorGrid\", [\"positions\", \"colors\"])\n",
    "\n",
    "\n",
    "def build_color_grid(x_count=10, y_count=10, z_count=10, twist=0):\n",
    "    \"\"\"\n",
    "    Create a cube of points with colors.\n",
    "\n",
    "    The total point cloud will have x_count * y_count * z_count points.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x_count, y_count, z_count:\n",
    "        Number of points in each dimension.\n",
    "    twist:\n",
    "        Angle to twist from bottom to top of the cube\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    grid = np.mgrid[\n",
    "        slice(-x_count, x_count, x_count * 1j),\n",
    "        slice(-y_count, y_count, y_count * 1j),\n",
    "        slice(-z_count, z_count, z_count * 1j),\n",
    "    ]\n",
    "\n",
    "    angle = np.linspace(-float(twist) / 2, float(twist) / 2, z_count)\n",
    "    for z in range(z_count):\n",
    "        xv, yv, zv = grid[:, :, :, z]\n",
    "        rot_xv = xv * cos(angle[z]) - yv * sin(angle[z])\n",
    "        rot_yv = xv * sin(angle[z]) + yv * cos(angle[z])\n",
    "        grid[:, :, :, z] = [rot_xv, rot_yv, zv]\n",
    "\n",
    "    positions = np.vstack([xyz.ravel() for xyz in grid])\n",
    "\n",
    "    colors = np.vstack(\n",
    "        [\n",
    "            xyz.ravel()\n",
    "            for xyz in np.mgrid[\n",
    "                slice(0, 255, x_count * 1j),\n",
    "                slice(0, 255, y_count * 1j),\n",
    "                slice(0, 255, z_count * 1j),\n",
    "            ]\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    return ColorGrid(positions.T, colors.T.astype(np.uint8))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04c095ef",
   "metadata": {},
   "source": [
    "## Logging some data\n",
    "\n",
    "Now we can log some data and add it to the recording, and display it using `notebook_show`.\n",
    "\n",
    "Note that displaying a recording will consume the data, so it will not be available for use in a subsequent cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92871ea1",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr.init(\"rerun_example_cube\")\n",
    "\n",
    "STEPS = 100\n",
    "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
    "for t in range(STEPS):\n",
    "    rr.set_time_sequence(\"step\", t)\n",
    "    cube = build_color_grid(10, 10, 10, twist=twists[t])\n",
    "    rr.log(\"cube\", rr.Points3D(cube.positions, colors=cube.colors, radii=0.5))\n",
    "\n",
    "rr.notebook_show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "187430e1",
   "metadata": {},
   "source": [
    "## Logging live data\n",
    "\n",
    "Using `rr.notebook_show` like above will buffer the data and only display it when the cell finishes running.\n",
    "\n",
    "If you want to see the data live, you can create a `Viewer` in one cell and then log data in another cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "052c4c6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr.init(\"rerun_example_cube\")\n",
    "viewer = rr.notebook.Viewer(display=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "246c8eb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import sleep\n",
    "\n",
    "STEPS = 100\n",
    "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
    "for t in range(STEPS):\n",
    "    sleep(0.1)\n",
    "    rr.set_time_sequence(\"step\", t)\n",
    "    cube = build_color_grid(10, 10, 10, twist=twists[t])\n",
    "    rr.log(\"cube\", rr.Points3D(cube.positions, colors=cube.colors, radii=0.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31d392a8",
   "metadata": {},
   "source": [
    "## Adjusting the view\n",
    "\n",
    "The  `notebook_show` method also lets you adjust properties such as width and height."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a1b0f66-4287-4705-8be5-ae837ffe3f90",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "small_cube = build_color_grid(3, 3, 3, twist=0)\n",
    "rr.log(\"small_cube\", rr.Points3D(small_cube.positions, colors=small_cube.colors, radii=0.5))\n",
    "\n",
    "rr.notebook_show(width=400, height=400)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36f9f61b",
   "metadata": {},
   "source": [
    "## Starting a new recording\n",
    "\n",
    "You can always start another recording by calling `rr.init(...)` again to reset the global stream, or alternatively creating a separate recording stream using `rr.new_recording` (discussed more below)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4cc33fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr.init(\"rerun_example_cube\")\n",
    "\n",
    "STEPS = 100\n",
    "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
    "for t in range(STEPS):\n",
    "    rr.set_time_sequence(\"step\", t)\n",
    "    h_grid = build_color_grid(10, 3, 3, twist=twists[t])\n",
    "    rr.log(\"h_grid\", rr.Points3D(h_grid.positions, colors=h_grid.colors, radii=0.5))\n",
    "    v_grid = build_color_grid(3, 3, 10, twist=twists[t])\n",
    "    rr.log(\"v_grid\", rr.Points3D(v_grid.positions, colors=v_grid.colors, radii=0.5))\n",
    "\n",
    "rr.notebook_show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9812634-067f-4e07-95fb-cb9a506c42d3",
   "metadata": {},
   "source": [
    "## Using blueprints\n",
    "\n",
    "Rerun blueprints can be used with `rr.notebook_show()`\n",
    "\n",
    "For example, we can split the two grids into their own respective space-views."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb8f7701",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr.init(\"rerun_example_cube\")\n",
    "\n",
    "blueprint = rrb.Blueprint(\n",
    "    rrb.Horizontal(\n",
    "        rrb.Spatial3DView(name=\"Horizontal grid\", origin=\"h_grid\"),\n",
    "        rrb.Spatial3DView(name=\"Vertical grid\", origin=\"v_grid\"),\n",
    "        column_shares=[2,1]),\n",
    "    collapse_panels=True\n",
    ")\n",
    "\n",
    "STEPS = 100\n",
    "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
    "for t in range(STEPS):\n",
    "    rr.set_time_sequence(\"step\", t)\n",
    "    h_grid = build_color_grid(10, 3, 3, twist=twists[t])\n",
    "    rr.log(\"h_grid\", rr.Points3D(h_grid.positions, colors=h_grid.colors, radii=0.5))\n",
    "    v_grid = build_color_grid(3, 3, 10, twist=twists[t])\n",
    "    rr.log(\"v_grid\", rr.Points3D(v_grid.positions, colors=v_grid.colors, radii=0.5))\n",
    "\n",
    "rr.notebook_show(blueprint=blueprint)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a63b30d7",
   "metadata": {},
   "source": [
    "## Extra convenience\n",
    "\n",
    "Rerun blueprints types also implement `_ipython_display_()` directly, so if a blueprint is the last element in your cell the right thing will happen.\n",
    "\n",
    "Note that this mechanism only works when you are using the global recording stream."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d1bf2ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr.init(\"rerun_example_cube\")\n",
    "\n",
    "STEPS = 100\n",
    "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
    "for t in range(STEPS):\n",
    "    rr.set_time_sequence(\"step\", t)\n",
    "    h_grid = build_color_grid(10, 3, 3, twist=twists[t])\n",
    "    rr.log(\"h_grid\", rr.Points3D(h_grid.positions, colors=h_grid.colors, radii=0.5))\n",
    "    v_grid = build_color_grid(3, 3, 10, twist=twists[t])\n",
    "    rr.log(\"v_grid\", rr.Points3D(v_grid.positions, colors=v_grid.colors, radii=0.5))\n",
    "\n",
    "rrb.Spatial3DView(name=\"Horizontal grid\", origin=\"h_grid\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef9087de-d090-4c90-ab3e-18c20c92bff6",
   "metadata": {},
   "source": [
    "## Working with non-global streams\n",
    "\n",
    "Sometimes it can be more explicit to work with specific (non-global recording) streams via the `new_recording` method.\n",
    "\n",
    "In this case, remember to call `notebook_show` directly on the recording stream. As noted above, there is no way to use a bare Blueprint object in conjunction with a non-global recording."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed86cc19-45cf-4c21-9a94-c4ce2ade7f65",
   "metadata": {},
   "outputs": [],
   "source": [
    "rec = rr.new_recording(\"rerun_example_cube_flat\")\n",
    "\n",
    "flat_grid = build_color_grid(20, 20, 1, twist=0)\n",
    "rec.log(\"flat_grid\", rr.Points3D(flat_grid.positions, colors=flat_grid.colors, radii=0.5))\n",
    "\n",
    "bp = rrb.Blueprint(collapse_panels=True)\n",
    "\n",
    "rec.notebook_show(blueprint=bp)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
